;; Multiarray type

;; by Konrad Hinsen
;; last updated June 5, 2010

;; Copyright (c) Konrad Hinsen, 2009-2010. All rights reserved.  The use
;; and distribution terms for this software are covered by the Eclipse
;; Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
;; which can be found in the file epl-v10.html at the root of this
;; distribution.  By using this software in any fashion, you are
;; agreeing to be bound by the terms of this license.  You must not
;; remove this notice, or any other, from this software.

(use 'nstools.ns)
(ns+ multiarray.type
  (:clone multiarray.utility)
  (:use multiarray.protocol)
  (:require [multiarray.descriptor :as descr]
            [multiarray.seqs-and-iterators :as si])
  (:from methods-a-la-carte.core deftype+ defimpl defcode expand-if)
  (:from methods-a-la-carte.implementations metadata keyword-lookup))

(set! *warn-on-reflection* true)

;;
;; Multiarray type check and type selection
;;

(defvar- multiarray-types #{})

(defn multiarray?
  [x]
  (if (get multiarray-types (type x))
    true
    false))

(defvar- multiarray-ctor-immutable {})
(defvar- multiarray-ctor-mutable {})

;
; Printing - the current code is simple and useful for debugging,
; but not for use in applications.
;

(defn- print-multiarray
  [#^String typename ma #^java.io.Writer w]
  (.write w "#:")
  (.write w typename)
  (.write w "{:shape ")
  (print-method (shape ma) w)
  (.write w ", :stride ")
  (print-method (vec (. (:descr ma) stride)) w)
  (.write w ", :offset ")
  (print-method (. (:descr ma) offset) w)
  (.write w ", :data ")
  (print-method (vec (take 10 (:data ma))) w)
  (.write w "}"))

;
; Macro for defining a multiarray type
;

(declare make-ma)

(defimpl java-interfaces []

  Object
    (equals [this o]
      (and (multiarray? o)
	   (descr/shape= descr (:descr o))
	   (all-true? (map = (elements this) (elements o)))))
    (hashCode [this]
      (reduce (fn [#^Integer h #^Object o] (+ (* 31 h) (.hashCode o)))
	      1 (elements this)))

  clojure.lang.Counted
    (count [this]
      (descr/assert-rank-not-0 descr)
      (aget #^ints (. descr shape) 0))

  clojure.lang.Indexed
    (nth [this i]
      (descr/assert-rank-not-0 descr)
      (slice this 0 i))

  clojure.lang.Sequential

  clojure.lang.Seqable
    (seq [this]
      (descr/assert-rank-not-0 descr)
      (let [n (aget #^ints (. descr shape) 0)]
	(map (partial slice this 0) (range 0 n))))

  java.lang.Iterable
    (iterator [this]
      (descr/assert-rank-not-0 descr)
      (let [n (aget #^ints (. descr shape) 0)]
	(si/make-dim0-iterator 0 n (partial slice this 0)))))

(defcode new-ma [mutable?]
  ~(expand-if ~mutable?
     (new ~this-type d data owner)
     (new ~this-type d data)))

(defimpl multiarray-methods [array-type mutable?]
  MultiArrayProtocol
    (rank [this]
      (descr/rank descr))
    (shape [this]
      (descr/shape descr))
    (elements [this]
      (flatten this))
    (reshape [this new-shape]
      (if-let [d (descr/reshape descr new-shape)]
	~(new-ma ~mutable?)
	~(expand-if ~mutable?
	   (throw (Exception. "can't reshape"))
	   (if (zero? (apply * (descr/shape descr)))
	     (throw (Exception.
		     "can't reshape empty to non-empty array"))
	     (let [ma (make-ma (element-class this) new-shape false)]
	       (acopy ~array-type (:data ma)
		      (cycle (elements this)))
	       ma)))))
    (sample [this dim start end step]
      (let [d (descr/sample descr dim start end step)]
	~(new-ma ~mutable?)))
    (slice [this dim index]
	   (let [d (descr/slice descr dim index)]
	     ~(expand-if ~mutable?
	        (new ~this-type d data owner)
		(if (zero? (descr/rank d))
		  (aget data (:offset d))
		  (new ~this-type d data)))))
    (transpose [this]
      (transpose this (range (dec (descr/rank descr)) -1 -1)))
    (transpose [this dim-permutation]
      (let [d (descr/transpose descr dim-permutation)]
	~(new-ma ~mutable?))))

;;;;;;;;;;

(deftype+ IntegerMultiArray
  [#^multiarray.descriptor.MultiArrayDescriptor descr
   #^ints data
   __meta]
  ~@(metadata __meta)
  ~@(keyword-lookup descr data)
  ~@(java-interfaces)
  ~@(multiarray-methods ints false))

(def multiarray-types
     (conj multiarray-types IntegerMultiArray))

(def multiarray-ctor-immutable
     (conj multiarray-ctor-immutable
	   [Integer/TYPE
	    (fn [descr array]
	      (new IntegerMultiArray descr array))]))

(defmethod print-method IntegerMultiArray [ma w]
  (print-multiarray (str (quote IntegerMultiArray)) ma w))

;;;;;;;;;;

(comment

(defmacro defmultiarray
  [typename array-type element-type mutable?]
  (template [typename     typename
	     fields       (let [fields [#^multiarray.descriptor.MultiArrayDescriptor 'descr
					(with-meta 'data {:tag array-type})]]
			    (if mutable?
			      (conj fields #^clojure.lang.Atom 'owner)
			      fields))
	     new-ma       (if mutable?
			    (template [typename typename]
			      (new ~typename d data owner))
			    (template [typename typename]
			      (new ~typename d data)))
	     new-ma-view  (if mutable?
			    (template [typename typename]
			      (new ~typename d data owner))
			    (template [typename typename]
			      (if (zero? (descr/rank d))
				(aget data (:offset d))
				(new ~typename d data))))
	     new-reshape  (if mutable?
			    `(throw (Exception. "can't reshape"))
			    `(if (zero? (apply * (descr/shape ~'descr)))
			       (throw (Exception.
				      "can't reshape empty to non-empty array"))
			       (let [~'ma (make-ma (element-class ~'this)
						   ~'new-shape false)]
				 (acopy ~array-type (:data ~'ma)
					(cycle (elements ~'this)))
				 ~'ma)))
	     type-kw      (keyword (str *ns*) (str typename))
	     element-type element-type
	     ctor-map     (if mutable?
			    'multiarray-ctor-mutable
			    'multiarray-ctor-immutable)]
    (do

      (deftype ~typename

	~fields

	Object
	  (equals [this o]
	    (and (multiarray? o)
		 (descr/shape= descr (:descr o))
		 (all-true? (map = (elements this) (elements o)))))
	  (hashCode [this]
	    (reduce (fn [#^Integer h #^Object o] (+ (* 31 h) (.hashCode o)))
		    1 (elements this)))

	clojure.lang.Counted
	  (count [this]
	    (descr/assert-rank-not-0 descr)
	    (aget #^ints (:shape descr) 0))

	clojure.lang.Indexed
	  (nth [this i]
	    (descr/assert-rank-not-0 descr)
	    (slice this 0 i))

	clojure.lang.Sequential

	clojure.lang.Seqable
	  (seq [this]
	    (descr/assert-rank-not-0 descr)
	    (let [n (aget #^ints (:shape descr) 0)]
	      (map (partial slice this 0) (range 0 n))))

	java.lang.Iterable
	  (iterator [this]
	    (descr/assert-rank-not-0 descr)
	    (let [n (aget #^ints (:shape descr) 0)]
	      (si/make-dim0-iterator 0 n (partial slice this 0))))

	MultiArrayLowLevelProtocol
	  (element-class [this]
	    (. (class data) getComponentType))
	  (element-array [this]
	    ; Case of non-simple arrays not handled very efficiently
	    (if (descr/simple? descr (alength data))
	      data
	      ((:element-array multiarray-lowlevel-method-defaults) this)))

	MultiArrayProtocol
	  (rank [this]
	    (descr/rank descr))
	  (shape [this]
            (descr/shape descr))
	  (elements [this]
	    (flatten this))
	  (reshape [this new-shape]
	    (if-let [d (descr/reshape descr new-shape)]
	      ~new-ma
	      ~new-reshape))
	  (sample [this dim start end step]
            (let [d (descr/sample descr dim start end step)]
	      ~new-ma))
	  (slice [this dim index]
            (let [d (descr/slice descr dim index)]
	      ~new-ma-view))
          (transpose [this]
	    (transpose this (range (dec (descr/rank descr)) -1 -1)))
	  (transpose [this dim-permutation]
            (let [d (descr/transpose descr dim-permutation)]
	      ~new-ma)))

      (def multiarray-types
	   (conj multiarray-types ~type-kw))

      (def ~ctor-map
	    (conj ~ctor-map [~element-type ~typename]))

      (defmethod print-method ~type-kw [ma w]
	(print-multiarray (str (quote ~typename)) ma w)))))


(defmultiarray ByteMultiArray    bytes   Byte/TYPE    false)
(defmultiarray ShortMultiArray   shorts  Short/TYPE   false)
(defmultiarray IntegerMultiArray ints    Integer/TYPE false)
(defmultiarray FloatMultiArray   floats  Float/TYPE   false)
(defmultiarray DoubleMultiArray  doubles Double/TYPE  false)
(defmultiarray MultiArray        objects Object       false)

(defmultiarray MutableByteMultiArray    bytes   Byte/TYPE    true)
(defmultiarray MutableShortMultiArray   shorts  Short/TYPE   true)
(defmultiarray MutableIntegerMultiArray ints    Integer/TYPE true)
(defmultiarray MutableFloatMultiArray   floats  Float/TYPE   true)
(defmultiarray MutableDoubleMultiArray  doubles Double/TYPE  true)
(defmultiarray MutableMultiArray        objects Object       true)
)

;;
;; Creation of new multiarrays
;;

(defn- make-ma
  [element-class shape mutable?]
  (let [descr      (descr/make-descr shape)
	data-array (make-array element-class (apply * shape))]
    (if mutable?
      (let [ctor (multiarray-ctor-mutable element-class)]
	(ctor descr data-array (atom (java.lang.Thread/currentThread))))
      (let [ctor (multiarray-ctor-immutable element-class)]
	(ctor descr data-array)))))

;;
;; Finalization of mutable arrays
;;

(defn finalize
  [ma]
  (when-not (compare-and-set! (:owner ma) (java.lang.Thread/currentThread) nil)
    (throw (Exception. "finalization allowed only by owner thread")))
  (let [ctor (multiarray-ctor-immutable (element-class ma))]
    (ctor (:descr ma) (:data ma))))

;;
;; Array constructors
;;

(set! *warn-on-reflection* false)

(defn zeros
  ([shape]
     (zeros shape Integer/TYPE false))
  ([shape element-class]
     (zeros shape element-class false))
  ([shape element-class mutable?]
     (let [ma (make-ma element-class shape mutable?)]
       (afill (:data ma) dummy 0)
       ma)))

(defn ma-range
  "Returns a multiarray whose elements are those of range."
  ([end]
     (ma-range 0 end 1))
  ([start end]
     (ma-range start end 1))
  ([start end step]
     (let [el-class (common-element-class [start end step])
	   cast     (cast-fn el-class)]
       (let [n (count (range start end step))
	     ma (make-ma el-class [n] false)]
	 (loop [i (int 0)
		s (range start end step)]
	   (when (seq s)
	     (do (aset (:data ma) i (cast (first s)))
		 (recur (inc i) (rest s)))))
	 ma))))

(set! *warn-on-reflection* true)

(defn into-ma
  ([object] (into-ma object false))
  ([object mutable?]
     (let [descr      (descr/make-descr (shape object))
	   data-array (element-array object)]
       (if mutable?
	 (let [ctor (multiarray-ctor-mutable (element-class object))]
	   (ctor descr data-array (atom (java.lang.Thread/currentThread))))
	 (let [ctor (multiarray-ctor-immutable (element-class object))]
	   (ctor descr data-array))))))


(set! *warn-on-reflection* false)

